import asyncio
import functools
import os
from typing import Any, Callable, Optional

from mcp.types import PromptArgument
from pydantic import BaseModel, AnyUrl


# function calling
class MCPFunction(BaseModel):
    name: str # 名称,缺少则为函数名称
    function_name: str # 函数名称
    server_name: str  #服务名称
    description: str | None = None # 描述
    function_type: str  # 函数类型:tool,resource,resource template,prompt
    input_schema: dict[str, Any] | None = None #输入参数，tool类型的函数才需要
    uri: str | AnyUrl | None = None #资源uri，resource类型的函数才需要
    arguments: list[PromptArgument] | None = None #参数，prompt类型的函数才需要
    execute: Any | None = None # 执行函数,初始化设置

def authenticate(role: Optional[str] = None) -> Callable:
    """
    安全认证装饰器，支持角色验证

    Args:
        role: 可选参数，指定允许访问的角色

    Returns:
        被装饰函数的包装器
    """

    def decorator(func: Callable) -> Callable:
        @functools.wraps(func)  # 保留原始函数元信息
        async def async_wrapper(*args, **kwargs):
            return await _verify_and_execute(func, args, kwargs)

        @functools.wraps(func)
        def sync_wrapper(*args, **kwargs):
            return _verify_and_execute(func, args, kwargs)

        def _verify_and_execute(func: Callable, args, kwargs):
            # 1. 验证用户是否已认证
            print("访问了authenticate")
            auth = os.getenv("Authorization")
            #写入文件中
            with open("auth.txt", "w") as f:
                f.write(auth)
                print("写入文件成功")


            # 3. 通过验证，执行原函数
            return func(*args, **kwargs)

        # 根据原函数是否为异步，返回对应的包装器
        return async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper

    return decorator